module fluid
use technical
use grid
use mod_2dflu
use fft
implicit none
public
save
double precision, pointer :: pomega(:,:),pDomega(:,:)
double precision, pointer :: pux(:,:),puy(:,:),ppsi(:,:)
double precision, pointer :: pomega_ux(:,:),pomega_uy(:,:)
contains
!******************************************
subroutine point_fluid
!
! Sets up the pointers for the fluid variables. 
!
! First omega which is solved for
  iomega=1
  pomega=>f(:,:,iomega)
  pDomega=>df(:,:,iomega)
! 
! Then the auxilliary variables.
!
  iux=2
  pux=>df(:,:,iux)
  iuy=3
  puy=>df(:,:,iuy)
  ipsi=4
  ppsi=>df(:,:,ipsi)
!  
! The two other arrays (omega ux, omega uy ) are stored in the
! same location as (ux, uy)
!
  pomega_ux=>df(:,:,iux)
  pomega_uy=>df(:,:,iuy)
!
endsubroutine point_fluid
!******************************************
subroutine initcond_fluid
!
! Sets up initial condition for velocity 
!
  select case (init_uu)
  case ('zero')
    call announce('init_uu:','zero, velocity and vorticity')
    pomega = 0.
  case ('white-omega')
    call announce('init_uu:','vorticity is white noise in space')
    call white_omega ()
  case default
    call fatal_error('init_uu:','initial condition for velocity not found')
  endselect
  call  announce('initcond_fluid:','Initial condition for fluid is set')
!
endsubroutine initcond_fluid
!**************************************************************!
subroutine initialize_diagnostic_fluid (n_var_print)
  integer :: n_var_print
  integer :: iprint
!
!n_var_print is the number of print.in quantities that are
! requested. 
!
  do iprint=1,n_var_print
!     call parse_name(iprint,c_vol_avg(iprint),cform(iname),'ekin',idiag_ekin)
  enddo
endsubroutine initialize_diagnostic_fluid
!**************************************************************!
subroutine get_df_fluid_part1 (istep,ireal,iimag,i2,k1,k2)
!
!Calculates the RHS of the hydrodynamic equations. 
!
  integer :: istep
  integer :: ireal,iimag,i2,k1,k2
  double precision :: rk2,nu,muk
!
! calculate psi and velocities. 
!
  call get_psivel (ireal,iimag,i2,k1,k2)
!
! calculate diagnostic (in Fourier space) if this is the first step of the multistep integrator
! There will be other conditions, e.g., diagnostic will not be calculated at every
! step etc. 
!
  if((istep.eq.1).and.ldiagnos) & 
       call diagnostic_fluid_fourier (ireal,iimag,i2,k1,k2)
!
!Then add contribution from viscosity (vis), hyperviscosity (vis2) and friction (mu)
! (If an exponential integrator is being used a different method has to be applied. )
!
!
  if(.not.(lexponential_integrator)) then
    if((vis.ne.0.).and.(vis2.ne.0.)) &
         call add_viscous_contrib (ireal,iimag,i2,k1,k2)
 else
 endif
!
endsubroutine get_df_fluid_part1
!***********************************************************!
subroutine fft_fluid_fourier_to_real
  integer :: i1,i2
!
! Call (inverse) FFT to transform omega and the velocities to real space
!
  call fft_inplace_inverse_2d(pomega,n1,n2)
  call fft_inplace_inverse_2d(pux,n1,n2)
  call fft_inplace_inverse_2d(puy,n1,n2)
!
!It is not necessary to transform psi to real space to time step our equation
!however it may be useful for diagnostic. Hence this is done, unless
!explicitly forbidden to save calculation time.
!
  if (linvert_psi) call fft_inplace_inverse_2d(ppsi,n1,n2)!
! -------normalize etc ------------------------------
!
  do i1=n1+1,n1+2
    pux(i1,:) = 0.0d0
    puy(i1,:) = 0.0d0
    pomega(i1,:) = 0.0d0
  enddo
  if (lun_normalized_fft) then
  pomega = pomega*scale
  pux = pux*scale; puy = puy*scale
  if (linvert_psi) ppsi=ppsi*scale
  else
  endif
!
!Part 1 is finished. Now all quantities are in real space. 
! pomega = omega (real space)
! pux = ux (real space) = (pomega_ux)
! puy = uy (real space) = (pomega_uy)
! ppsi = psi (real space) [linvert_psi=T]
! ppsi = psi (Fourier space) [linvert_psi=F]
!
endsubroutine fft_fluid_fourier_to_real
!***********************************************************!
subroutine df_fluid_real_space (istep,i1,i2)
!
!Calculates quantities in real space needed to construct df and
! also real space diagnostics
!
  integer :: istep,i1,i2
!
  if((istep.eq.1).and.(ldiagnos)) call diagnostic_fluid_real (i1,i2)
!
!construct ux omega and uy omega
!
! The eqn we are solving is \partial_t \omega + J = 0 
! where J = div(u\omega)  In Fourier space this eqn will look like
! \partial_t \omega = - i k FFT(u\omega)
! So we first construct u\omega is real space, which is a vector. 
! We store in the same array as ux and uy
!
  pomega_ux(i1,i2) = pomega(i1,i2)*pux(i1,i2)
  pomega_uy(i1,i2) = pomega(i1,i2)*pux(i1,i2)
!
endsubroutine df_fluid_real_space 
!***********************************************************!
subroutine fft_fluid_real_to_fourier
!
! Now FFT to fourier space to obtain FFT(u\omega) which is a 
! two dimensional vector array, or two two-dimensional arrays. 
!
  call fft_inplace_forward_2d(pomega_ux,n1,n2)
  call fft_inplace_forward_2d(pomega_uy,n1,n2)
!
endsubroutine fft_fluid_real_to_fourier
!***********************************************************!
subroutine get_df_fluid_part2 (istep,ireal,iimag,i2,k1,k2)
!
!Calculates the RHS of the hydrodynamic equations. (part 2) 
!
  integer :: istep
  integer :: i2,ireal,iimag,k1,k2
  integer :: ksqr,kx,ky
  double precision :: rk2
  double complex :: Jacobean_complex,omega_ux_complex,omega_uy_complex
!
! Now we are back in Fourier space. 
! (except psi which is in Real space / Fourier depending on 
! linvert_psi variable is true / false
! and omega which is in Real space  ) 
! pomega_ux = FFT(omega ux) (= pux )
! pomega_uy = FFT(omega uy) (= puy )
! The last equality, inside parenthesis, is because the two pointer
! arrays, pomega_ux and pux point to the same memory location. 
! pomega = omega (in real space) 
! (same is true for pomega_uy and puy)
! Ready to add the contribution from the Jacobean.
!
  if (.not.lexponential_integrator) then
     ksqr = k1*k1 + k2*k2 
     rk2 = two_pi_byL*two_pi_byL*dfloat(ksqr)
     kx = two_pi_byL*dfloat(k1)
     ky = two_pi_byL*dfloat(k2)
     if((ksqr.gt.kasqr).or.(ksqr.eq.0))then
        pDomega(ireal,i2) = 0.
        pDomega(iimag,i2) = 0.
     else
        omega_ux_complex=dcmplx(pomega_ux(ireal,i2),pomega_ux(iimag,i2))
        omega_uy_complex=dcmplx(pomega_uy(ireal,i2),pomega_uy(iimag,i2))
        Jacobean_complex=zi*(kx*omega_ux_complex+ky*omega_uy_complex)
        pDomega(ireal,i2) = pDomega(ireal,i2) - real(Jacobean_complex)
        pDomega(iimag,i2) = pDomega(iimag,i2) - aimag(Jacobean_complex)
     endif
  else
  endif
!
endsubroutine get_df_fluid_part2
!******************************************
subroutine get_psivel (ireal,iimag,i2,k1,k2)
!
! Obtain ux, uy and psi in Fourier space from omega
!
  integer :: ireal,iimag,i2,k1,k2
  integer :: ksqr,kx,ky
  double precision :: rk2
!-------------------------
  ksqr = k1*k1 + k2*k2 
  rk2 = two_pi_byL*two_pi_byL*dfloat(ksqr)
  kx = two_pi_byL*dfloat(k1)
  ky = two_pi_byL*dfloat(k2)
  if((ksqr.gt.kasqr).or.(ksqr.eq.0))then
     pomega(ireal,i2) = 0.
     pomega(iimag,i2) = 0.
     ppsi(ireal,i2) = 0.
     ppsi(iimag,i2) = 0.
     pux(ireal,i2) = 0.
     pux(iimag,i2) = 0.
     puy(ireal,i2) = 0.
     puy(iimag,i2) = 0.
  else
     ppsi(ireal,i2) = -pomega(ireal,i2)/rk2
     ppsi(iimag,i2) = -pomega(iimag,i2)/rk2
     pux(ireal,i2) = -ky*ppsi(iimag,i2)
     pux(iimag,i2) = ky*ppsi(ireal,i2)
     puy(ireal,i2) = kx*ppsi(iimag,i2)
     puy(iimag,i2) = -kx*ppsi(ireal,i2)
  endif
!--------------------- 
endsubroutine get_psivel
!******************************************!
subroutine diagnostic_fluid_fourier (ireal,iimag,i2,k1,k2)
!
!Calculates diagnostic equatities for the fluid variables in Fourier space.
!
  integer :: ireal,iimag,i2,k1,k2
  double precision :: rk2,rk,energy,dissipation,enstrophy,entdiss
  integer :: ksqr,mshl 
!--------------------------------------------------------------
  ksqr = k1*k1 + k2*k2  
  rk2 = two_pi_byL*two_pi_byL*dfloat(ksqr)
  rk = dsqrt(dble(ksqr))
  mshl = nint(rk)
  if(ksqr.gt.kasqr)then
     E_Omega(mshl,1) = 0.
     E_Omega(mshl,2) = 0.
     E_Omega(mshl,3) = 0.
  else
     energy = (one_by_nsqr**2)*(& 
          pux(ireal,i2)**2 +  pux(iimag,i2)**2 & 
          + puy(ireal,i2)**2 + puy(iimag,i2)**2 )
     enstrophy = (one_by_nsqr**2)*(&
          pomega(ireal,i2)**2 + pomega(iimag,i2)**2) 
     entdiss = rk2*enstrophy
     if((k1.eq.0).or.(k1.eq.n1h)) then
        E_Omega(mshl,1) = E_Omega(mshl,1) + energy
        E_Omega(mshl,2) = E_Omega(mshl,2) + enstrophy 
        E_Omega(mshl,3) = E_Omega(mshl,3) + entdiss
     else
        E_Omega(mshl,1) = E_Omega(mshl,1) + 2.0d0*energy
        E_Omega(mshl,2) = E_Omega(mshl,2) + 2.0d0*enstrophy
        E_Omega(mshl,3) = E_Omega(mshl,3) + 2.0d0*entdiss
     endif
  endif
!! --------------------------------------------------------
endsubroutine diagnostic_fluid_fourier
!******************************************
subroutine add_viscous_contrib (ireal,iimag,i2,k1,k2)
  integer :: ireal,iimag,i2,k1,k2
  integer :: ksqr
  double precision :: muk,nu,rk2
!
!Note on muk: We can cutoff the friction for higher k if we want. Or put any such
! function of ksqr. A mu which depends anisotropically on k (i.e., not on ksqr alone)
! is not inconceivable, although can in principle be incorporated. 
! A ksqr dependent mu can be obtained by multiplying mu by the relevant
! function of ksq to obtain muk below. 
!
  muk = mu
  if((ksqr.gt.kasqr).or.(ksqr.eq.0))then
     pDomega(ireal,i2) = 0.
     pDomega(iimag,i2) = 0. 
     pomega(ireal,i2)=0.
     pomega(iimag,i2) = 0.
  else
     rk2 = two_pi_byL*two_pi_byL*dfloat(ksqr)
     nu = vis + vis2*rk2
     pDomega(ireal,i2) = pDomega(ireal,i2) - (nu*rk2+muk)*pomega(ireal,i2)
     pDomega(iimag,i2) = pDomega(iimag,i2) - (nu*rk2+muk)*pomega(iimag,i2)
  endif
!
endsubroutine add_viscous_contrib
!******************************************
subroutine diagnostic_fluid_real (i1,i2)
!
  integer :: i1,i2
!
!calculates diagnostic for the fluid in real space 
!
endsubroutine diagnostic_fluid_real
!******************************************
subroutine white_omega
!
!vorticity is white noise with random phase
!
  double precision,dimension(3) :: ran
  double precision :: rk2,rk,ek1,vk,p11,p12,p31,p22,p33,p23,rk2inv
  double precision :: iniamp,omegak,n_sqr,rkini,mabs,prim,pert,x,y
  integer :: i1,i2,i3,k1,k2,k3,ksqr,mshl,ik,iseed,ireal,iimag,m1,m2
  integer,dimension(3) :: tarray
  integer,dimension(2) :: tofday
  integer :: ierr
!! -----------------------------------------------------------------
!! vorticity has a flat spectra with random phase 
  n_sqr = dfloat(n1*n2)
  iseed = 10;
!! --------- 
  do i2 = 1,n2
     y = i2*dx
     do i1 = 1,n1
        x = i1*dx
        prim = -4.0d0*(dcos(4.0d0*y) + dcos(4.0d0*x))
        pert = 0.0d0
        do m2 = 0,2!Nx/2
           do m1 = 0,2!Ny/2
              mabs = dsqrt(dble(m1*m1 + m2*m2))
!
              if(m1**2+m2**2.ne.0)then
                 pert = pert + &
                      1.0d-4*(dsin(m1*x+m2*y) + dcos(m1*x+m2*y))*dble(m2*m2)/mabs
              endif
	      pert = 0.d0;
!
           enddo
        enddo
        pomega(i1,i2) = (prim + pert)*vis
     enddo
  enddo
!  
      call fft_inplace_forward_2d(pomega,n1,n2)
!
endsubroutine white_omega
!******************************************
end module fluid
